import numpy as np
import python.version as version
from python.helpers import (
    mod_inverse,
    reduce_sym_list,
)


def merge_lwe(lwe_instance):
    # 把间隔为256的项合并
    n = version.KYBER_N
    a = lwe_instance.a
    b = lwe_instance.b
    q = lwe_instance.q
    coff_a = [[]] * n
    coff_b = [0] * n

    for i in range(n):
        coff_a[i] = [sum(items) for items in zip(*[a[i + k * n] for k in range(version.KYBER_k)])]
        coff_b[i] = (sum(b[i + k * n] for k in range(version.KYBER_k))) % q

    return coff_a, coff_b


def load_coff_s(a, b, guessed_e, guessed_s, know_indices):
    m = version.KYBER_k * version.KYBER_N
    n = version.KYBER_N
    q = 3329

    coff_a_s = []
    coff_b_s = []

    for i in range(n):
        if i in know_indices:
            b_tmp = b[i]
            a_tmp = []
            # b_i-e_i
            b_tmp = (b_tmp - guessed_e[i]) % q

            # A_i_j * s_j
            for j in range(0, m):
                if j + n in know_indices:
                    # 减去系数
                    b_tmp = (b_tmp - a[i][j] * guessed_s[j]) % q
                else:
                    # 添加系数
                    a_tmp.append(a[i][j])
            coff_a_s.append(a_tmp)
            coff_b_s.append(b_tmp)

    coff_a_s = np.array(coff_a_s)
    coff_b_s = np.array(coff_b_s)

    return coff_a_s, coff_b_s


def gaussian_elimination(a, b, q=3329):
    a = a.astype(int)
    b = b.astype(int)

    # 合并 A 和 b
    ab = np.column_stack((a, b))
    dim_m, dim_n = ab.shape

    # 初等变换
    for i in range(min(dim_m, dim_n)):
        # 检查除数是否为0
        if ab[i, i] == 0:
            # 寻找该元素不为0的行
            for k in range(i + 1, dim_m):
                if ab[k, i] != 0:
                    ab[[i, k]] = ab[[k, i]]  # 交换行 numpy的advance index，ab[[k,i]]即选择第k行和第i行
                    break
            else:
                # 该列全0 跳过
                continue

        # 当前行归一
        inv = mod_inverse(ab[i, i])
        ab[i] = (ab[i] * inv) % q

        # 用当前行消去下面的行
        inv = mod_inverse(ab[i, i])
        for j in range(i + 1, dim_m):
            factor = (ab[j, i] * inv) % q
            ab[j] = (ab[j] - factor * ab[i]) % q

    return ab[:, :-1], ab[:, -1]


def solve_s(a, b, q=3329):
    # 求解s
    dim = len(a)
    solution = [0 for _ in range(dim)]

    # 从最后一行倒往前消
    for i in range(dim - 1, -1, -1):
        solution[i] = b[i]
        for j in range(i + 1, dim):
            solution[i] = (solution[i] - (a[i, j] * solution[j])) % q
        solution[i] = (solution[i] * mod_inverse(a[i, i])) % q

    return reduce_sym_list(solution)


def solve_e(a, b, known_indices, guessed_s):
    q = 3329
    n = version.KYBER_N
    m = version.KYBER_k * version.KYBER_N
    solution = []

    # e_i 系数为1 直接赋值就行
    for i in range(n):
        if i not in known_indices:
            b_tmp = b[i]
            for j in range(0, m):
                b_tmp = (b_tmp - a[i][j] * guessed_s[j]) % q
            solution.append(b_tmp)

    return reduce_sym_list(solution)


def solve_equation(propagation_data, step, step_rank):
    key = propagation_data.key
    n = version.KYBER_k * version.KYBER_N
    if step == -1:
        # no use_best_step
        current_step = sorted(propagation_data.steps.items(), key=lambda x: x[0])[step][
            1
        ]
    else:
        current_step = propagation_data.steps[step]
        # current_step is PropagationDataStep instance
    recovered_num = current_step.recovered_coefficients
    ordered_key_indices = current_step.ordered_key_indices
    known_indices = ordered_key_indices[:recovered_num]
    guessed_e = [current_step.guessed_key[i] for i in range(0, version.KYBER_N)]
    guessed_s = [current_step.guessed_key[i] for i in range(version.KYBER_N, n + version.KYBER_N)]
    assert all(
        (
            known_e_i == propagation_data.key[i]
            for i, known_e_i in enumerate(guessed_e)
            if i in known_indices
        )
    )
    assert all(
        (
            known_s_i == propagation_data.key[i + version.KYBER_N]
            for i, known_s_i in enumerate(guessed_s)
            if i + version.KYBER_N in known_indices
        )
    )
    unknown_indices = sorted(
        ordered_key_indices[recovered_num:]
    )  # 未知系数系数 按照索引大小排序
    unknown_e = [i for i in unknown_indices if i < version.KYBER_N]  # 未知e索引
    unknown_s = [i for i in unknown_indices if i >= version.KYBER_N]  # 未知s索引

    # 合并矩阵
    merge_lwe_a, merge_lwe_b = merge_lwe(propagation_data.lwe_instance)

    # 求解e_i已知的行
    coff_A_s, coff_B_s = load_coff_s(
        merge_lwe_a, merge_lwe_b, guessed_e, guessed_s, known_indices
    )
    new_A, new_B = gaussian_elimination(coff_A_s, coff_B_s)

    # 检查消元完是不是方阵 检查后面的行有没有非0的
    last_nonzero_row = np.max(np.where(np.any(new_A != 0, axis=1)))
    assert np.all(new_B[last_nonzero_row + 1:] == 0)
    assert np.all(new_A[last_nonzero_row + 1:, :] == 0)
    new_A = new_A[0: last_nonzero_row + 1]
    new_B = new_B[0: last_nonzero_row + 1]
    assert new_A.ndim == 2 and new_A.shape[0] == new_A.shape[1]

    # 求解 赋值 检查
    solution_s = solve_s(new_A, new_B)
    for i in range(len(solution_s)):
        index_s = unknown_s[i]
        guessed_s[index_s - version.KYBER_N] = solution_s[i]
    assert guessed_s == key[version.KYBER_N:]

    # 再算e_i未知的行 直接代入就行
    solution_e = solve_e(merge_lwe_a, merge_lwe_b, known_indices, guessed_s)
    for i in range(len(solution_e)):
        index_e = unknown_e[i]
        guessed_e[index_e] = solution_e[i]
    assert guessed_e == key[0: version.KYBER_N]
    assert guessed_e + guessed_s == key
    propagation_data.solution = guessed_e + guessed_s
    with open("output.txt", "a") as file:
        file.write(f"Recovered all {len(key)} coefficients (e, s).\n")
    propagation_data.set_lattice_data(None, (0, 0), 0, step, step_rank, 0, 0)

    # 要不要算origin_e

    return True
